package com.chj.zhongji.class09.test01;

//https://blog.csdn.net/clover_my/article/details/105906522
//https://blog.csdn.net/ZhEnG_jIe_YuAn/article/details/107302128
public class Code02_KthSmallest {
	public static int kthSmallest(int[][] mat, int k) {
		int m = mat.length;
		int[] res = mat[0];
		for (int i = 1; i < m; i++) {
			res = topKSum(res, mat[i], k);
		}
		return res[k - 1];
	}

	public static class HeapNode {
		public int row;
		public int col;
		public int value;

		public HeapNode(int row, int col, int value) {
			this.row = row;
			this.col = col;
			this.value = value;
		}
	}

	public static int[] topKSum(int[] a1, int[] a2, int topK) {
		if (a1 == null || a2 == null || topK < 1) {
			return null;
		}

		topK = Math.min(topK, a1.length * a2.length);
		HeapNode[] heap = new HeapNode[topK + 1];
		int heapSize = 0;
		int rowl = a1.length - 1;
		int coll = a2.length - 1;
		int headR = -1;
		int headC = -1;
		int uR = -1;
		int uC = -1;
		int lR = -1;
		int lC = -1;
		heapInsert(heap, heapSize, 0, 0, a1[0] + a2[0]);
		heapSize++;
		boolean[][] set = new boolean[a1.length][a2.length];
		int[] res = new int[topK];
		int resIndex = 0;
		while (resIndex != topK) {
			HeapNode head = popHead(heap, heapSize);
			heapSize--;
			res[resIndex++] = head.value;

			headR = head.row;
			headC = head.col;

			uR = headR + 1;
			uC = headC;
			if (uR <= rowl && !set[uR][uC]) {
				heapInsert(heap, heapSize, uR, uC, a1[uR] + a2[uC]);
				heapSize++;
				set[uR][uC] = true;
			}

			lR = headR;
			lC = headC + 1;
			if (lC <= coll && !set[lR][lC]) {
				heapInsert(heap, heapSize, lR, lC, a1[lR] + a2[lC]);
				heapSize++;
				set[lR][lC] = true;
			}
		}
		return res;
	}

	public static HeapNode popHead(HeapNode[] heap, int heapSize) {
		HeapNode res = heap[0];
		swap(heap, 0, heapSize - 1);
		heapSize--;
		heap[heapSize] = null;
		heapify(heap, 0, heapSize);
		return res;
	}

	public static void heapInsert(HeapNode[] heap, int index, int row, int col, int value) {
		heap[index] = new HeapNode(row, col, value);
		int parent = (index - 1) / 2;
		while (heap[index].value < heap[parent].value) {
			swap(heap, index, parent);
			index = parent;
			parent = (index - 1) / 2;
		}
	}

	public static void heapify(HeapNode[] heap, int index, int heapSize) {
		int left = index * 2 + 1;
		int right = left;
		int largest = index;
		while (left < heapSize) {

			if (right < heapSize && heap[right].value < heap[left].value) {
				largest = right;
			} else {
				largest = left;
			}

			largest = heap[index].value > heap[largest].value ? largest : index;

			if (largest == index) {
				break;
			}

			swap(heap, largest, index);
			index = largest;
			left = index * 2 + 1;
			right = left + 1;
		}
	}

	public static void swap(HeapNode[] heap, int index1, int index2) {
		HeapNode tmp = heap[index1];
		heap[index1] = heap[index2];
		heap[index2] = tmp;
	}

	public static void main(String[] args) {
		int[][] matrix = new int[][] { { 1, 3, 11 }, { 2, 4, 6 } };
		System.out.println(kthSmallest(matrix, 5));

		int[][] matrix2 = new int[][] { { 1, 3, 11 }, { 2, 4, 6 } };
		System.out.println(kthSmallest(matrix2, 9));

		int[][] matrix3 = new int[][] { { 1, 10, 10 }, { 1, 4, 5 }, { 2, 3, 6 } };
		System.out.println(kthSmallest(matrix3, 7));

		int[][] matrix4 = new int[][] { { 1, 1, 10 }, { 2, 2, 9 } };
		System.out.println(kthSmallest(matrix4, 7));
	}
}
